#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward A. Roualdes
#         [2014.05.29]
#         University of Kentucky

from __future__ import print_function
import utils
import sys
import re
import os.path
from collections import defaultdict
import subprocess

def mkPhase(outdir, cleanSNP = None):
    u"""create PHASE file from inferred and aligned haplotypes stored within outdir"""

    print("Making PHASE file...")
    bases = ['A', 'C', 'G', 'T', 'a', 'c', 'g', 't']
    other = ['-', 'n', 'N']
    allbases = bases+other

    # make PHASE directory
    phaseDir = os.path.join(outdir, 'phase')
    utils.mkdir(phaseDir, overwrite=True)

    # get aligned haplotypes data
    data = utils.files2dict(os.path.join(outdir, 'haplotypes'), '.*\\.aligned$')

    ## for each .*\\.aligned file's matrix of data, find and store snps
    for d, indiv, locus in data:
        # initialize empty containers for positions, nucleotides, and locus types
        pos = [] # list of position(i) indicating the position of polymorphic nucleotide
        loc_type = []               # list of locus types indicating biallelic or microsatellite
        snp = defaultdict(list) # dict of {key = indiv_name1 : [nucleotide_pos[1], ..., nucleotide_pos[2]]}

        # if data exists, read it (file by file)
        assert d, "no data in file %s.fna.aligned." % locus
        read_mat, indiv_names = utils.read_matrix(d)

        # search each row at column j for snp
        for j in xrange(len(read_mat[0])):
            col = [row[j] for row in read_mat if row[j] in allbases]
            ucol = utils.unique(col)
            if len([x for x in ucol if x in bases]) > 1: # polymorphic position?

                # store position number
                pos.append(str(j+1))
                
                # store locus type
                # [2014.05.29] TODO: surely we need better checking for biallelic vs micosatellite
                # & implement cleanSNP option
                if len(ucol) == 2: 
                    lt = 'S'
                else: 
                    lt = 'M'
                loc_type.append(lt)

                # store nucleotide in map keyed by individual name
                for i in xrange(len(col)):
                    ind_name = indiv_names[i]
                    if col[i] in other: # catch missing values
                        if lt == 'S':
                            snp[ind_name].append('?')
                        else:
                            snp[ind_name].append('-1')
                    else:
                        snp[ind_name].append( utils.let2num(col[i]) )

                
        # write info to file
        if pos:
            with open(os.path.join(phaseDir, '%s.phase.inp' % locus), 'w') as f:
                print(len(read_mat)/2, file=f) # number of individuals
                print(len(pos), file=f)      # number of loci
                print('P ' + ' '.join(pos), file=f) # positions
                print(' '.join(loc_type), file=f) # locus types
                used = []
                for name in snp.keys():
                    if name[:-2] not in used:
                        idx = name.rfind('_') # index of last underscore in name, it should be created by NextAllele
                        print(name[:idx], file=f) # individual name, no individual number: _1 or _2
                        print(' '.join(snp[name]), file=f) # individual nucleotides, for first individual (of two)
                        jdx = [x[:-2] for x in snp.keys() if x != name].index(name[:idx]) # index of key of second individual (of 2)
                        name2 = snp.keys()[jdx]
                        print(' '.join(snp[name2]), file=f) # nucleotides for second individual
                        used = used + [name[:-2]] # ensure we don't repeat individuals


    return

    
def files2phase(phaseDirec):
    u"""creates iterable of files stored in outdir with extension .phase.inp"""

    l = []
    pat = re.compile(".*\\.phase.inp$")
    # scan files in directory
    for root, dirs, files in os.walk(phaseDirec):
        for fname in filter(pat.match, files):
            l.append(os.path.join(root, fname))
    return l


def runPhase(f, phasePath, outdir):
    u"""runs program PHASE after creating all PHASE files"""
    infile = f
    outfile = os.path.splitext(f)[0] + '.out'
    # [2014.05.29] TODO: add PHASE flags & try to caputure output (cause otherwise its a mess)
    subprocess.call([phasePath, infile, outfile])

    return
    
    
